import logging
import functools
import json
from typing import List, Optional
import lru

from web3 import Web3
from web3.types import (
    Wei,
    ChecksumAddress,
)
from web3 import middleware
from web3.middleware import validation

from coin.exchange.uniswap_base.types import AddressLike
from coin.exchange.uniswap_base.exceptions import (
    InvalidToken,
    InvalidPool,
)
from coin.exchange.uniswap_base.util import (
    str_to_addr,
    addr_to_str,
    load_contract,
    is_same_address,
)
from coin.exchange.bsc_base.token import (
    BEP20Token,
    Pair,
)
from coin.exchange.bsc_base.constants import (
    WBNB_ADDRESS,
    SYMBOL_TO_ADDRESS,
    TOKEN_TO_ADDRESS,
)
from coin.exchange.uniswap_base.constants import (
    netid_to_name,
    UNISWAP_V2_FACTORY_CONTRACT_ADDRESS,
    UNISWAP_V2_ROUTER_CONTRACT_ADDRESS,
)

def estimate_quote_qty(quote_address):
  # Assume 5000 USD worth of quote. If quote is BNB, then worth of 10 BNB.
  if quote_address in (TOKEN_TO_ADDRESS['USDT'], TOKEN_TO_ADDRESS['BUSD'], TOKEN_TO_ADDRESS['USDC']):
    return 5000
  elif quote_address == WBNB_ADDRESS:
    return 10
  elif quote_address == TOKEN_TO_ADDRESS['BTC']:
    return 0.1
  elif quote_address == TOKEN_TO_ADDRESS['ETH']:
    return 1
  else:
    raise NotImplementedError(quote_address)


class PancakeswapManager:
  def __init__(
      self,
      provider: str = None,
      web3: Web3 = None,
      version: int = 2,
      logger=None,
  ) -> None:
    """
    :param address: The public address of the ETH wallet to use.
    :param provider: Web3 provider URI.
    :param web3: Can be optionally set to a custom Web3 instance.
    """
    self._version = version
    if web3 is not None:
      self._w3 = web3
    else:
      self._provider = provider
      self._w3 = Web3(Web3.HTTPProvider(self._provider, request_kwargs={"timeout": 60}))
      self._w3.middleware_onion.add(middleware.time_based_cache_middleware)
      latest_block_based_cache_middleware = middleware.construct_latest_block_based_cache_middleware(
        cache_class=functools.partial(lru.LRU, 256),
        rpc_whitelist=middleware.cache.BLOCK_NUMBER_RPC_WHITELIST,
        default_average_block_time=3,
      )
      self._w3.middleware_onion.add(latest_block_based_cache_middleware)
      self._w3.middleware_onion.add(middleware.simple_cache_middleware)
      self._w3.middleware_onion.inject(middleware.geth_poa_middleware, layer=0)
      validation.METHODS_TO_VALIDATE = []
    netid = int(self._w3.net.version)
    netname = netid_to_name[netid]
    self._logger = logger or logging.getLogger('PancakeswapManager')
    self._logger.info(f"Web3 client {self._w3} connected. netname={netname}")
    self._last_bid0_prices = {}
    self._last_ask0_prices = {}

    if self._version == 2:
      router_contract_addr = UNISWAP_V2_ROUTER_CONTRACT_ADDRESS[netname]
      self._router_contract_address: AddressLike = str_to_addr(router_contract_addr)

      factory_contract_addr = UNISWAP_V2_FACTORY_CONTRACT_ADDRESS[netname]
      self._factory_contract = load_contract(
          self._w3,
          abi_path="python/coin/exchange/pancakeswap_v2/abi/factory.abi",
          address=str_to_addr(factory_contract_addr),
      )
      self._router_contract = load_contract(
          self._w3,
          abi_path="python/coin/exchange/pancakeswap_v2/abi/router.abi",
          address=self._router_contract_address,
      )
      self._pairs = {}
      self._pair_event_filters = {}
      self._tx_hash_set = set()
    else:
      raise ValueError(f"Invalid version: {self._version}, only 2 supported!")

  @functools.lru_cache()
  def _get_pair(self, symbol: str) -> Pair:
    """
    Retrieves metadata from the Pair contract,
    """
    pair_contract_addr = SYMBOL_TO_ADDRESS[symbol]
    pair_contract = load_contract(
        self._w3,
        abi_path="python/coin/exchange/pancakeswap_v2/abi/pair.abi",
        address=pair_contract_addr,
    )
    try:
      token0 = self.get_token(pair_contract.functions.token0().call())
      token1 = self.get_token(pair_contract.functions.token1().call())
      base, quote = symbol.split('-')
      if base in token0.symbol.upper():
        base_then_quote = True
      else:
        base_then_quote = False
    except Exception as e:
      self._logger.error(
          f"Exception occurred while trying to get pair {symbol}: {e}")
      raise InvalidPool(pair_contract_addr)
    return Pair(symbol, pair_contract, token0, token1, 3000, base_then_quote)

  @functools.lru_cache()
  def _get_pair_filter(self, symbol: str):
    pair = self._get_pair(symbol)
    self._pairs[symbol] = pair
    self._pair_event_filters[symbol] = self._pairs[symbol].contract.events.Swap \
        .create_filter(fromBlock='latest')
    return self._pair_event_filters[symbol]

  def get_as_bid0(self, symbol: str, fee=3000) -> list:
    pair = self._get_pair(symbol)
    if pair.base_then_quote:
      base_address = pair.token0.address
      base_decimals = pair.token0.decimals
      quote_address = pair.token1.address
      quote_decimals = pair.token1.decimals
    else:
      base_address = pair.token1.address
      base_decimals = pair.token1.decimals
      quote_address = pair.token0.address
      quote_decimals = pair.token0.decimals
    quote_qty = estimate_quote_qty(quote_address)
    bid_qty = self.get_price_output(
        base_address, quote_address, int(quote_qty * 10**quote_decimals), fee=pair.fee)
    bid_qty /= 10**base_decimals
    bid_price = quote_qty / bid_qty
    self._last_bid0_prices[symbol] = bid_price
    return [bid_price, bid_qty]

  def get_as_ask0(self, symbol: str, fee=500) -> list:
    pair = self._get_pair(symbol)
    if pair.base_then_quote:
      base_address = pair.token0.address
      base_decimals = pair.token0.decimals
      quote_address = pair.token1.address
      quote_decimals = pair.token1.decimals
    else:
      base_address = pair.token1.address
      base_decimals = pair.token1.decimals
      quote_address = pair.token0.address
      quote_decimals = pair.token0.decimals
    quote_qty = estimate_quote_qty(quote_address)
    ask_qty = self.get_price_input(
        quote_address, base_address, int(quote_qty * 10**quote_decimals), fee=pair.fee)
    ask_qty /= 10**base_decimals
    ask_price = quote_qty / ask_qty
    self._last_ask0_prices[symbol] = ask_price
    return [ask_price, ask_qty]

  def get_trades(self, symbol: str) -> list:
    pair = self._get_pair(symbol)
    filter = self._get_pair_filter(symbol)
    try:
      events = filter.get_new_entries()
    except ValueError as e:
      self._logger.error(e)
      # Try to reset filter.
      self._get_pair_filter.cache_clear()
      filter = self._get_pair_filter(symbol)
      events = filter.get_all_entries()
    event_jsons = []
    if len(events) >= 1:
      tx_hash = events[0]['transactionHash']
      if tx_hash not in self._tx_hash_set:
        self._tx_hash_set.add(tx_hash)
      else:
        return []

      last_bid0_price = self._last_bid0_prices.get(symbol, None)
      if last_bid0_price is not None:
        valid_price_floor = last_bid0_price * (1 - 0.2)
      else:
        # bbo price is broken, then the subscriber is damaged.
        # So just publish whatever fetched from trade if possible.
        valid_price_floor = 0

      last_ask0_price = self._last_ask0_prices.get(symbol, None)
      if last_ask0_price is not None:
        valid_price_ceil = last_ask0_price * (1 + 0.2)
      else:
        # bbo price is broken, then the subscriber is damaged.
        # So just publish whatever fetched from trade if possible.
        valid_price_ceil = float('inf')

      for swap_event in events:
        swap_event_json = json.loads(Web3.to_json(swap_event))
        amount0 = swap_event_json['args']['amount0Out']
        amount1 = swap_event_json['args']['amount1In']
        amount0_out = True
        if amount0 == 0 and amount1 == 0:
          amount0 = swap_event_json['args']['amount0In']
          amount1 = swap_event_json['args']['amount1Out']
          amount0_out = False
        if pair.base_then_quote:
          base_amount = amount0 / (10**pair.token0.decimals)
          quote_amount = amount1 / (10**pair.token1.decimals)
          if amount0_out:
            base_out = True
          else:
            base_out = False
        else:
          base_amount = amount1 / (10**pair.token1.decimals)
          quote_amount = amount0 / (10**pair.token0.decimals)
          if amount0_out:
            base_out = False
          else:
            base_out = True
        if base_out:
          # Contract perspective is opposite to trader perspective.
          side = 'B'
        else:
          side = 'S'
        price = quote_amount / base_amount
        if price > valid_price_floor and price < valid_price_ceil:
          event_jsons.append({
              'price': price,
              'qty': base_amount,
              'side': side,
          })
        else:
          self._logger.error(f"Filter out invalid price {symbol}: {price}")
    return event_jsons

  def get_reserves(self, symbol: str) -> list:
    pair = self._get_pair(symbol)

    if self._version == 2:
      reserve0, reserve1, _timestamp = pair.contract.functions.getReserves().call()
      reserve0 = reserve0 / (10 ** pair.token0.decimals)
      reserve1 = reserve1 / (10 ** pair.token1.decimals)
      if pair.base_then_quote:
        return [reserve0, reserve1]
      else:
        return [reserve1, reserve0]
    else:
      raise ValueError(self._version)

  def get_price_input(
      self,
      token0: AddressLike,
      token1: AddressLike,
      qty: int,
      fee: int = None,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """Returns the amount of the input token you get for `qty` of the output token"""
    if fee is None:
      fee = 3000  # 3000/1e6 => 0.3%

    if token0 == WBNB_ADDRESS:
      return self._get_bnb_token_input_price(token1, Wei(qty), fee)
    elif token1 == WBNB_ADDRESS:
      return self._get_token_bnb_input_price(token0, qty, fee)
    else:
      return self._get_token_token_input_price(token0, token1, qty, fee, route)

  def get_price_output(
      self,
      token0: AddressLike,
      token1: AddressLike,
      qty: int,
      fee: int = None,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """Returns the amount of input token you need to get `qty` of the output token"""
    if fee is None:
      fee = 3000  # 3000/1e6 => 0.3%

    if is_same_address(token0, WBNB_ADDRESS):
      return self._get_bnb_token_output_price(token1, qty, fee)
    elif is_same_address(token1, WBNB_ADDRESS):
      return self._get_token_bnb_output_price(token0, Wei(qty), fee)
    else:
      return self._get_token_token_output_price(token0, token1, qty, fee, route)

  def _get_bnb_token_input_price(self, token: AddressLike, qty: Wei, fee: int) -> Wei:
    """Public price for BNB to Token trades with an exact input."""
    if self._version == 2:
      price = self._router_contract.functions.getAmountsOut(
          qty, [WBNB_ADDRESS, token]).call()[-1]
    else:
      raise ValueError(self._version)
    return price

  def _get_token_bnb_input_price(self, token: AddressLike, qty: int, fee: int) -> int:
    """Public price for token to BNB trades with an exact input."""
    if self._version == 2:
      price = self._router_contract.functions.getAmountsOut(
          qty, [token, WBNB_ADDRESS]).call()[-1]
    else:
      raise ValueError(self._version)
    return price

  def _get_bnb_token_output_price(self, token: AddressLike, qty: int, fee: int = None) -> Wei:
    """Public price for ETH to Token trades with an exact output."""
    if self._version == 2:
      route = [WBNB_ADDRESS, token]
      price = self._router_contract.functions.getAmountsIn(qty, route).call()[0]
    else:
      raise ValueError(self._version)
    return price

  def _get_token_bnb_output_price(self, token: AddressLike, qty: Wei, fee: int = None) -> int:
    """Public price for token to ETH trades with an exact output."""
    if self._version == 2:
      route = [token, WBNB_ADDRESS]
      price = self._router_contract.functions.getAmountsIn(qty, route).call()[0]
    else:
      raise ValueError(self._version)
    return price

  def _get_token_token_input_price(
      self,
      token0: AddressLike,  # input token
      token1: AddressLike,  # output token
      qty: int,
      fee: int,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """
    Public price (i.e. amount of output token received) for token to token trades with an exact input.

    :param fee: (v3 only) The pool's fee in hundredths of a bip, i.e. 1e-6 (3000 is 0.3%)
    """
    if route is None:
      route = [token0, WBNB_ADDRESS, token1]
      self._logger.debug(f"No route specified, assuming route: {route}")

    price: int = self._router_contract.functions.getAmountsOut(qty, route).call()[-1]
    return price

  def _get_token_token_output_price(
      self,
      token0: AddressLike,  # input token
      token1: AddressLike,  # output token
      qty: int,
      fee: int = None,
      route: Optional[List[AddressLike]] = None,
  ) -> int:
    """
    Public price (i.e. amount of input token needed) for token to token trades with an exact output.

    :param fee: (v3 only) The pool's fee in hundredths of a bip, i.e. 1e-6 (3000 is 0.3%)
    """
    if not route:
      route = [token0, WBNB_ADDRESS, token1]
      self._logger.debug(f"No route specified, assuming route: {route}")

    price: int = self._router_contract.functions.getAmountsIn(qty, route).call()[0]
    return price

  @functools.lru_cache()
  def get_token(self, address: AddressLike) -> BEP20Token:
    """
    Retrieves metadata from the BEP20 contract of a given token,
    like its name, symbol, and decimals.
    """
    token_contract = load_contract(
        self._w3,
        abi_path="python/coin/exchange/bsc_base/abi/bep20.abi",
        address=address)
    try:
      name = token_contract.functions.name().call()
      symbol = token_contract.functions.symbol().call()
      decimals = token_contract.functions.decimals().call()
    except Exception as e:
      self._logger.error(
          f"Exception occurred while trying to get token {addr_to_str(address)}: {e}")
      raise InvalidToken(address)
    return BEP20Token(symbol, address, name, decimals)
